// CryptProtectDataTest.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"

void PrintHelp()
{
	cerr << "CryptProtectDataTest.exe - A Utility to test Windows CryptProtectData functionality" << endl;
	cerr << "Version " << MYVERSION << " - azengel@pgp.com" << endl;
	cerr << "Usage:" << endl;
	cerr << "CryptProtectDataTest.exe --encrypt <TEXT>     : Encrypt the text specified in <TEXT> and write to a file" << endl;
	cerr << "                                                Remember to use quotes (\") when specifying a text with spaces!" << endl;
	cerr << "CryptProtectDataTest.exe --decrypt <FILENAME> : Decrypt the data in file <FILENAME>, show contents and write to a file" << endl << endl;
	cerr << "Notes:" << endl;
	cerr << "The output filenames are automatically generated with a timestamp in the filename and will be generated in the" << endl;
	cerr << "current working directory." << endl;
	cerr << "When encrypting: The text will be written to a file with name '<TIMESTAMP>_encrypted.txt', also the plain text will be" << endl;
	cerr << "                 written to a file named '<TIMESTAMP>_plaintext.txt' for reference." << endl;
	cerr << "When decrypting: The decrypted content will be printed on screen and will be written to a file"<<endl;
	cerr << "                 named '<INPUTFILE>_decrypt.txt'" << endl;
	cerr << "" << endl;
}

void PrintError(LPTSTR lpszFunction) 
{ 
    // Retrieve the system error message for the last-error code

    LPVOID lpMsgBuf;
    LPVOID lpDisplayBuf;
    DWORD dw = GetLastError(); 

    FormatMessage(
        FORMAT_MESSAGE_ALLOCATE_BUFFER | 
        FORMAT_MESSAGE_FROM_SYSTEM |
        FORMAT_MESSAGE_IGNORE_INSERTS,
        NULL,
        dw,
        MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
        (LPTSTR) &lpMsgBuf,
        0, NULL );

    // Display the error message and exit the process

    lpDisplayBuf = (LPVOID)LocalAlloc(LMEM_ZEROINIT, 
        (lstrlen((LPCTSTR)lpMsgBuf) + lstrlen((LPCTSTR)lpszFunction) + 40) * sizeof(TCHAR)); 
    StringCchPrintf((LPTSTR)lpDisplayBuf, 
        LocalSize(lpDisplayBuf) / sizeof(TCHAR),
        TEXT("%s failed with error %d: %s"), 
        lpszFunction, dw, lpMsgBuf); 
	printf("%s",lpDisplayBuf);
   // MessageBox(NULL, (LPCTSTR)lpDisplayBuf, TEXT("Error"), MB_OK); 

    LocalFree(lpMsgBuf);
    LocalFree(lpDisplayBuf);
}

int WriteToFile(string fileName, string content)
{
	try 
	{
		ofstream plaintextfile=ofstream(fileName,ofstream::binary);
		plaintextfile.write (content.c_str(),content.length());
		plaintextfile.close();
	} catch (exception e)
	{
		printf("Could not write to file '%s'\n", fileName.c_str());
		return 1;
	}
	return 0;
}

int EncryptStringToFile(string plainText,string fileName)
{
	DATA_BLOB DataIn;
	DATA_BLOB DataOut;
	BYTE *pbDataInput = (BYTE*)plainText.c_str();
	DWORD cbDataInput = strlen((char *)pbDataInput)+1;
	int ret = 0;

	printf("Encrypting data '%s'\n",plainText.c_str());

	//--------------------------------------------------------------------
	// Initialize the DataIn structure.
	DataIn.pbData = pbDataInput;    
	DataIn.cbData = cbDataInput;

	//--------------------------------------------------------------------
	//  Begin protect phase. Note that the encryption key is created
	//  by the function and is not passed.
	BOOL result = false;
	result = CryptProtectData(
		 &DataIn,
		 NULL,								 // A description string
											 // to be included with the
											 // encrypted data. 
		 NULL,                               // Optional entropy not used.
		 NULL,                               // Reserved.
		 NULL,                               // Pass NULL for the 
											 // prompt structure.
		 CRYPTPROTECT_UI_FORBIDDEN,
		 &DataOut);
	if(result)
	{
		 printf("The encryption phase worked.\n");
		 try
		 {
				ofstream encryptedfile=ofstream(fileName,ofstream::binary);
				encryptedfile.write ((char*)DataOut.pbData,DataOut.cbData);
				encryptedfile.close();
		 } catch (exception e)
		 {
				printf ("Could not write to file '%s'\n", fileName.c_str());
				ret = 1;
		 }
			printf ("Wrote encrypted data to file '%s'\n", fileName.c_str());
			ret = 0;
	}
	else
	{
		PrintError(TEXT("CryptProtectData"));
		printf("Encryption error using CryptProtectData.\n");
		ret = 2;
	}
	
	LocalFree(&DataOut);
	LocalFree(&DataIn);
	return ret;
}

int DecryptFileToFile(string inFile, string outFile)
{
	// Decrypt data from DATA_BLOB DataOut to DATA_BLOB DataVerify.
	//--------------------------------------------------------------------
	// Declare and initialize variables.
	//--------------------------------------------------------------------
	DATA_BLOB DataIn;
	DATA_BLOB DataOut;
	BYTE *pbDataInput;
	DWORD cbDataInput;
	int ret = 0;

	printf ("Decrypting data in file '%s'\n",inFile.c_str());

	//--------------------------------------------------------------------
	try 
	{
		ifstream infile=ifstream(inFile,ifstream::binary);
		// get size of file
		infile.seekg(0,ifstream::end);
		cbDataInput=(long)infile.tellg();
		infile.seekg(0);

		// initialize buffer
		pbDataInput = new BYTE [cbDataInput];

		// read content of infile
		infile.read((char*)pbDataInput,cbDataInput);
		infile.close();
	} catch (exception e)
	{
		printf("Could not read from file '%s'",inFile.c_str());
		ret = 5;
		goto done;
	}

	// Initialize the DataIn structure.
	DataIn.pbData = pbDataInput;    
	DataIn.cbData = cbDataInput;
	//--------------------------------------------------------------------
	//   Begin unprotect phase.

	if (CryptUnprotectData(
			&DataIn,
			NULL,
			NULL,                 // Optional entropy
			NULL,                 // Reserved
			NULL,                 // Here, the optional 
								  // prompt structure is not
								  // used.
			CRYPTPROTECT_UI_FORBIDDEN,
			&DataOut))
	{
		printf("The decrypted data is: '%s'\n", DataOut.pbData);
		try
		{
			ofstream decryptedfile=ofstream(outFile,ofstream::binary);
			decryptedfile.write ((char*)DataOut.pbData,DataOut.cbData-1);
			decryptedfile.close();
			printf("Wrote decrypted data to file '%s'\n",outFile.c_str());
		} catch (exception e)
		{
			printf ("Could not write to file ''\n",outFile.c_str());
			ret = 7;
			goto done;
		}
	}
	else
	{
		PrintError(TEXT("CryptUnProtectData"));
		printf("Decryption error!");
		ret = 6;
		goto done;
	}

	

	done:
	LocalFree(&DataOut);
	LocalFree(&DataIn);
	return ret;
}

int _tmain(int argc, char* argv[])
{
	int res = 0;
	string argFlag;

	//handle command line input
	if (argc < 3 || argc >3) {
		PrintHelp();
		res=99;
		return(res);
	}

	// get timestampt for filenames
	std::time_t t = std::time(0);
	std::stringstream ss;
	ss << t;
	std::string ts = ss.str();

	argFlag = string(argv[1]);
	if (!argFlag.compare("--encrypt"))
	{
		// write plaintext to file for reference
		// encrypt something
		stringstream fileNamePlain;
		fileNamePlain << ts << "_plaintext.txt";
		if (WriteToFile(fileNamePlain.str(),string(argv[2])))
		{
			res=77;
			goto done;
		}
		printf("Wrote plaintext data to file '%s' for reference\n",fileNamePlain.str().c_str());
		// encrypt something
		stringstream fileNameEnc;
		fileNameEnc << ts << "_encrypted.txt";
		res = EncryptStringToFile(string(argv[2]),fileNameEnc.str());

	} else if (!argFlag.compare("--decrypt"))
	{ 
		//decrypt a file
		string filename = string(argv[2]);
		filename.append("_decrypted.txt");
		res = DecryptFileToFile(string(argv[2]),filename);
	} else
	{
		PrintHelp();
		res=99;
		goto done;
	}

done:
	return res;

}

